{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from __future__ import print_function, absolute_import\n",
    "\n",
    "import os\n",
    "import numpy as np\n",
    "import json\n",
    "import random\n",
    "import math\n",
    "\n",
    "import torch\n",
    "import torch.utils.data as data\n",
    "\n",
    "from pose.utils.osutils import *\n",
    "from pose.utils.imutils import *\n",
    "from pose.utils.transforms import *\n",
    "from pose.utils.evaluation  import final_preds\n",
    "import pose.models as models\n",
    "\n",
    "import glob\n",
    "import cv2\n",
    "from tqdm import tqdm\n",
    "import imageio\n",
    "\n",
    "from scipy.io import loadmat\n",
    "import scipy.misc\n",
    "import scipy.ndimage\n",
    "\n",
    "import imgaug as ia\n",
    "import imgaug.augmenters as iaa\n",
    "from imgaug.augmentables.kps import KeypointsOnImage\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def cv2_plot_lines(frame, pts, order):\n",
    "    color_mapping = {1: [255,0,255], 2: [255,0,0], 3: [255,0,127], 4: [255,255,255], 5: [0,0,255],\n",
    "                 6: [0,127,255], 7: [0,255,255], 8: [0,255,0], 9: [200,162,200]} \n",
    "    point_size = 3\n",
    "    if order==0:\n",
    "        # other animals\n",
    "        # plot neck-eyes\n",
    "        cv2.line(frame, (pts[[0,2],:][0,0], pts[[0,2],:][0,1]), (pts[[0,2],:][1,0], pts[[0,2],:][1,1]), color_mapping[1], point_size)\n",
    "        cv2.line(frame, (pts[[1,2],:][0,0], pts[[1,2],:][0,1]), (pts[[1,2],:][1,0], pts[[1,2],:][1,1]), color_mapping[2], point_size)\n",
    "\n",
    "        # plot legs\n",
    "        cv2.line(frame, (pts[[3,8],:][0,0], pts[[3,8],:][0,1]), (pts[[3,8],:][1,0], pts[[3,8],:][1,1]), color_mapping[5], point_size)\n",
    "        cv2.line(frame, (pts[[8,14],:][0,0], pts[[8,14],:][0,1]), (pts[[8,14],:][1,0], pts[[8,14],:][1,1]), color_mapping[5], point_size)\n",
    "\n",
    "        cv2.line(frame, (pts[[4,9],:][0,0], pts[[4,9],:][0,1]), (pts[[4,9],:][1,0], pts[[4,9],:][1,1]), color_mapping[6], point_size)\n",
    "        cv2.line(frame, (pts[[9,15],:][0,0], pts[[9,15],:][0,1]), (pts[[9,15],:][1,0], pts[[9,15],:][1,1]), color_mapping[6], point_size)\n",
    "\n",
    "        cv2.line(frame, (pts[[5,10],:][0,0], pts[[5,10],:][0,1]), (pts[[5,10],:][1,0], pts[[5,10],:][1,1]), color_mapping[7], point_size)\n",
    "        cv2.line(frame, (pts[[10,16],:][0,0], pts[[10,16],:][0,1]), (pts[[10,16],:][1,0], pts[[10,16],:][1,1]), color_mapping[7], point_size)\n",
    "\n",
    "        cv2.line(frame, (pts[[6,11],:][0,0], pts[[6,11],:][0,1]), (pts[[6,11],:][1,0], pts[[6,11],:][1,1]), color_mapping[8], point_size)\n",
    "        cv2.line(frame, (pts[[11,17],:][0,0], pts[[11,17],:][0,1]), (pts[[11,17],:][1,0], pts[[11,17],:][1,1]), color_mapping[8], point_size)\n",
    "\n",
    "        # plot hip-necks\n",
    "        cv2.line(frame, (pts[[12,7],:][0,0], pts[[12,7],:][0,1]), (pts[[12,7],:][1,0], pts[[12,7],:][1,1]), color_mapping[1], point_size)\n",
    "        cv2.line(frame, (pts[[13,7],:][0,0], pts[[13,7],:][0,1]), (pts[[13,7],:][1,0], pts[[13,7],:][1,1]), color_mapping[2], point_size)\n",
    "    elif order==1:\n",
    "        # elephant\n",
    "        cv2.line(frame, (pts[[0,2],:][0,0], pts[[0,2],:][0,1]), (pts[[0,2],:][1,0], pts[[0,2],:][1,1]), color_mapping[1], point_size)\n",
    "        cv2.line(frame, (pts[[1,2],:][0,0], pts[[1,2],:][0,1]), (pts[[1,2],:][1,0], pts[[1,2],:][1,1]), color_mapping[2], point_size)\n",
    "\n",
    "        # plot legs\n",
    "        cv2.line(frame, (pts[[3,8],:][0,0], pts[[3,8],:][0,1]), (pts[[3,8],:][1,0], pts[[3,8],:][1,1]), color_mapping[5], point_size)\n",
    "        cv2.line(frame, (pts[[8,14],:][0,0], pts[[8,14],:][0,1]), (pts[[8,14],:][1,0], pts[[8,14],:][1,1]), color_mapping[5], point_size)\n",
    "\n",
    "        cv2.line(frame, (pts[[4,9],:][0,0], pts[[4,9],:][0,1]), (pts[[4,9],:][1,0], pts[[4,9],:][1,1]), color_mapping[6], point_size)\n",
    "        cv2.line(frame, (pts[[9,15],:][0,0], pts[[9,15],:][0,1]), (pts[[9,15],:][1,0], pts[[9,15],:][1,1]), color_mapping[6], point_size)\n",
    "\n",
    "        cv2.line(frame, (pts[[5,10],:][0,0], pts[[5,10],:][0,1]), (pts[[5,10],:][1,0], pts[[5,10],:][1,1]), color_mapping[7], point_size)\n",
    "        cv2.line(frame, (pts[[10,16],:][0,0], pts[[10,16],:][0,1]), (pts[[10,16],:][1,0], pts[[10,16],:][1,1]), color_mapping[7], point_size)\n",
    "\n",
    "        cv2.line(frame, (pts[[6,11],:][0,0], pts[[6,11],:][0,1]), (pts[[6,11],:][1,0], pts[[6,11],:][1,1]), color_mapping[8], point_size)\n",
    "        cv2.line(frame, (pts[[11,17],:][0,0], pts[[11,17],:][0,1]), (pts[[11,17],:][1,0], pts[[11,17],:][1,1]), color_mapping[8], point_size)\n",
    "\n",
    "        # plot hip-necks\n",
    "        cv2.line(frame, (pts[[12,7],:][0,0], pts[[12,7],:][0,1]), (pts[[12,7],:][1,0], pts[[12,7],:][1,1]), color_mapping[1], point_size)\n",
    "        cv2.line(frame, (pts[[13,7],:][0,0], pts[[13,7],:][0,1]), (pts[[13,7],:][1,0], pts[[13,7],:][1,1]), color_mapping[2], point_size)\n",
    "\n",
    "        \n",
    "        cv2.line(frame, (pts[[18,19],:][0,0], pts[[18,19],:][0,1]), (pts[[18,19],:][1,0], pts[[18,19],:][1,1]), color_mapping[1], point_size)\n",
    "        cv2.line(frame, (pts[[19,20],:][0,0], pts[[19,20],:][0,1]), (pts[[19,20],:][1,0], pts[[19,20],:][1,1]), color_mapping[1], point_size)\n",
    "        cv2.line(frame, (pts[[20,21],:][0,0], pts[[20,21],:][0,1]), (pts[[20,21],:][1,0], pts[[20,21],:][1,1]), color_mapping[1], point_size)\n",
    "        cv2.line(frame, (pts[[21,22],:][0,0], pts[[21,22],:][0,1]), (pts[[21,22],:][1,0], pts[[21,22],:][1,1]), color_mapping[1], point_size)\n",
    "        cv2.line(frame, (pts[[22,23],:][0,0], pts[[22,23],:][0,1]), (pts[[22,23],:][1,0], pts[[22,23],:][1,1]), color_mapping[1], point_size)\n",
    "        cv2.line(frame, (pts[[23,24],:][0,0], pts[[23,24],:][0,1]), (pts[[23,24],:][1,0], pts[[23,24],:][1,1]), color_mapping[1], point_size)\n",
    "\n",
    "\n",
    "def cv2_visualize_keypoints(frame, pts, num_pts=18, order=0):\n",
    "    \n",
    "    points = pts\n",
    "    x = []\n",
    "    y = []\n",
    "    for i in range(num_pts):\n",
    "        x.append(points[i][0])\n",
    "        y.append(points[i][1])\n",
    "        # plot keypoints on each image \n",
    "        cv2.circle(frame,(x[-1],y[-1]), 3, (0,255,0), -1)   \n",
    "    cv2_plot_lines(frame, points, order)\n",
    "    return frame\n",
    "\n",
    "def load_animal(data_dir='./', animal='horse'):\n",
    "    \"\"\"\n",
    "    Output:\n",
    "    img_list: Nx3   # each image is associated with a shot-id and a shot-id frame_id,\n",
    "                    # e.g. ('***.jpg', 100, 2) means the second frame in video_id 100.\n",
    "    anno_list: Nx3  # (x, y, visiblity)\n",
    "    \"\"\"\n",
    "\n",
    "    range_path = os.path.join(data_dir, 'behaviorDiscovery2.0/ranges', animal, 'ranges.mat')\n",
    "    landmark_path = os.path.join(data_dir, 'behaviorDiscovery2.0/landmarks', animal)\n",
    "\n",
    "    img_list = []  # img_list contains all image paths\n",
    "    anno_list = [] # anno_list contains all anno lists\n",
    "    range_file = loadmat(range_path)\n",
    "\n",
    "    for video in range_file['ranges']:\n",
    "        # range_file['ranges'] is a numpy array [Nx3]: shot_id, start_frame, end_frame\n",
    "        shot_id = video[0]\n",
    "        landmark_path_video = os.path.join(landmark_path, str(shot_id)+'.mat')\n",
    "\n",
    "        if not os.path.isfile(landmark_path_video):\n",
    "            continue\n",
    "        landmark_file = loadmat(landmark_path_video)\n",
    "\n",
    "        for frame in range(video[1], video[2]+1): # ??? video[2]+1\n",
    "            frame_id = frame - video[1]\n",
    "            img_name = '0'*(8-len(str(frame))) + str(frame) + '.jpg'\n",
    "            img_list.append([img_name, shot_id, frame_id])\n",
    "            \n",
    "            coord = landmark_file['landmarks'][frame_id][0][0][0][0]\n",
    "            vis = landmark_file['landmarks'][frame_id][0][0][0][1]\n",
    "            landmark = np.hstack((coord, vis))\n",
    "            anno_list.append(landmark[:18,:])\n",
    "            \n",
    "    return img_list, anno_list\n",
    "\n",
    "def dataset_filter(anno_list):\n",
    "    \"\"\"\n",
    "    output:\n",
    "    idxs: valid_idxs after filtering\n",
    "    \"\"\"\n",
    "    num_kpts = anno_list[0].shape[0]\n",
    "    idxs = []\n",
    "    for i in range(len(anno_list)):\n",
    "        s = sum(anno_list[i][:,2])\n",
    "        if s>num_kpts//2:\n",
    "            idxs.append(i)\n",
    "    return idxs\n",
    "\n",
    "def im_to_torch(img):\n",
    "    img = np.transpose(img, (2, 0, 1)) # C*H*W\n",
    "    img = to_torch(img).float()\n",
    "    if img.max() > 1:\n",
    "        img /= 255\n",
    "    return img\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Visualization for horses/tigers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "../checkpoint/synthetic_animal/horse/combineds5r5_decay0.01_parallel/model_best.pth.tar\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|          | 0/10152 [00:00<?, ?it/s]/media/jm/000C65DB000784DF/workspace/python37/lib/python3.7/site-packages/ipykernel_launcher.py:51: DeprecationWarning: `imread` is deprecated!\n",
      "`imread` is deprecated in SciPy 1.0.0, and will be removed in 1.2.0.\n",
      "Use ``imageio.imread`` instead.\n",
      "/media/jm/000C65DB000784DF/workspace/python37/lib/python3.7/site-packages/torch/nn/modules/upsampling.py:129: UserWarning: nn.Upsample is deprecated. Use nn.functional.interpolate instead.\n",
      "  warnings.warn(\"nn.{} is deprecated. Use nn.functional.interpolate instead.\".format(self.name))\n",
      "/media/jm/000C65DB000784DF/workspace/python37/lib/python3.7/site-packages/ipykernel_launcher.py:109: DeprecationWarning: `imresize` is deprecated!\n",
      "`imresize` is deprecated in SciPy 1.0.0, and will be removed in 1.3.0.\n",
      "Use Pillow instead: ``numpy.array(Image.fromarray(arr).resize())``.\n",
      "  0%|          | 44/10152 [00:11<43:37,  3.86it/s]"
     ]
    }
   ],
   "source": [
    "global_animal = 'horse' # can be horse, tiger\n",
    "order = 0 # 0: for other animals | 1: for elephant\n",
    "global_dataset = 'synthetic_animal' # models trained using synthetic datasets\n",
    "nParts = 18 # number of keypoints\n",
    "is_part = True # True: multitask setting | False: keypoint only\n",
    "\n",
    "\n",
    "if not os.path.exists(os.path.join('./demo', global_animal)):\n",
    "    os.makedirs(os.path.join('./demo', global_animal))\n",
    "\n",
    "# define directories\n",
    "img_folder = '/media/jm/000C65DB000784DF/workspace/animals/behaviorDiscovery/'\n",
    "img_list, anno_list = load_animal(data_dir=img_folder, animal=global_animal)\n",
    "img_idxs = dataset_filter(anno_list)\n",
    "\n",
    "# define the model\n",
    "if is_part:\n",
    "    # multi-task checkpoint\n",
    "    checkpoint_path = os.path.join(\"./checkpoint/\",global_dataset,global_animal+\"_multitask/model_best.pth.tar\")\n",
    "else:\n",
    "    # keypoint only checkpoint\n",
    "    checkpoint_path = os.path.join(\"./checkpoint/\",global_dataset,global_animal+'_ssl/synthetic_animal_sp.pth.tar\")\n",
    "print(checkpoint_path)\n",
    "meanstd_file = os.path.join('./data/', global_dataset, global_animal+'_combineds5r5_texture', 'mean.pth.tar')\n",
    "\n",
    "\n",
    "# Define the codec and create VideoWriter object\n",
    "fourcc = cv2.VideoWriter_fourcc(*'XVID')\n",
    "out = cv2.VideoWriter('output.avi',fourcc, 20.0, (512,256))\n",
    "\n",
    "# load model\n",
    "if is_part:\n",
    "    model1 = models.__dict__['hg_multitask'](num_stacks=4, num_blocks=1, num_classes=nParts, resnet_layers=50)\n",
    "else:\n",
    "    model1 = models.__dict__['hg'](num_stacks=4, num_blocks=1, num_classes=nParts, resnet_layers=50)\n",
    "model1 = torch.nn.DataParallel(model1).cuda()\n",
    "checkpoint = torch.load(checkpoint_path)\n",
    "model1.load_state_dict(checkpoint['state_dict'])\n",
    "idxs = np.arange(nParts)\n",
    "\n",
    "# calculate meand and std\n",
    "mean, std = torch.load(meanstd_file)['mean'], torch.load(meanstd_file)['std']\n",
    "\n",
    "for i in tqdm(range(len(img_idxs))):\n",
    "    # color_mapping for segmentations\n",
    "    # head, eye, ear, torso, left_front, right_front, left_back, right_back, tail\n",
    "    color_mapping = {1: [255,0,0], 2: [203,192,255], 3: [255,0,127], 4: [255,255,255], 5: [0,0,255],\n",
    "                 6: [0,127,255], 7: [0,255,255], 8: [0,255,0], 9: [200,162,200]}\n",
    "    \n",
    "    # calculate keypoints\n",
    "    img = scipy.misc.imread(os.path.join(img_folder, 'behaviorDiscovery2.0/', global_animal, img_list[img_idxs[i]][0]), mode='RGB')\n",
    "    img_path = img_list[img_idxs[i]][0]\n",
    "    frame = img.copy()\n",
    "    img = im_to_torch(img)\n",
    "    \n",
    "    # get correct scale and center\n",
    "    x_min = float(np.min(anno_list[img_idxs[i]][:,0] \\\n",
    "                         [anno_list[img_idxs[i]][:,0]>0]))\n",
    "    x_max = float(np.max(anno_list[img_idxs[i]][:,0] \\\n",
    "                         [anno_list[img_idxs[i]][:,0]>0]))\n",
    "    y_min = float(np.min(anno_list[img_idxs[i]][:,1] \\\n",
    "                         [anno_list[img_idxs[i]][:,1]>0]))\n",
    "    y_max = float(np.max(anno_list[img_idxs[i]][:,1] \\\n",
    "                         [anno_list[img_idxs[i]][:,1]>0]))\n",
    "\n",
    "    c = torch.Tensor(( (x_min+x_max)/2.0, (y_min+y_max)/2.0 ))\n",
    "    s = max(x_max-x_min, y_max-y_min)/200.0 * 1.5\n",
    "    rot = 0\n",
    "    \n",
    "    inp = crop(img, c, s, [256, 256], rot)\n",
    "    \n",
    "    frame = torch.Tensor(frame.transpose(2,0,1))\n",
    "    frame = crop(frame, c, s, [256, 256], rot)\n",
    "    frame = (frame.numpy().transpose(1,2,0))*255\n",
    "    frame = np.uint8(frame)\n",
    "    frame = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)\n",
    "#         frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)\n",
    "    \n",
    "    inp_show = inp.clone()\n",
    "    inp = color_normalize(inp, mean, std)\n",
    "\n",
    "    inp = inp.unsqueeze(0).cuda()\n",
    "    if is_part:\n",
    "        output, mask = model1(inp)\n",
    "    else:\n",
    "        output = model1(inp)\n",
    "    score_map = output[-1].cpu() if type(output) == list else output.cpu()\n",
    "\n",
    "    preds = final_preds(score_map, [c], [s], [64, 64])\n",
    "    preds = preds.squeeze(0)\n",
    "    pts = preds.clone().cpu().numpy()\n",
    "    \n",
    "    # get confidence score\n",
    "    confidence_score = np.max(score_map.detach().cpu().numpy(), axis=(0,2,3))\n",
    "    confidence = confidence_score>0.5\n",
    "    transformed_preds = np.zeros((nParts, 2), dtype=int)\n",
    "    for j in range(nParts):\n",
    "        transformed_preds[j] = transform(pts[j]+1, c, s, [256, 256], invert=0, rot=0)\n",
    "    \n",
    "    # get predicted mask\n",
    "    if is_part:\n",
    "        _, pred_seg = torch.max(mask[-1], 1)\n",
    "        part_seg = pred_seg.detach().cpu().numpy().transpose(1,2,0).astype(np.uint8).repeat(3, axis=2)\n",
    "        for i in range(1,10):\n",
    "            part_seg[:,:,0][part_seg[:,:,0]==i] = color_mapping[i][0]\n",
    "            part_seg[:,:,1][part_seg[:,:,1]==i] = color_mapping[i][1]    \n",
    "            part_seg[:,:,2][part_seg[:,:,2]==i] = color_mapping[i][2]\n",
    "\n",
    "        pred_mask = scipy.misc.imresize(part_seg, [256,256,3])\n",
    "        cv2.imshow('mask', pred_mask)\n",
    "\n",
    "#     # color_mapping for kpts\n",
    "#     color_mapping = {1: [255,0,255], 2: [255,0,0], 3: [255,0,127], 4: [255,255,255], 5: [0,0,255],\n",
    "#                  6: [0,127,255], 7: [0,255,255], 8: [0,255,0], 9: [200,162,200]}    \n",
    "    \n",
    "    # plot keypoints\n",
    "    frame = cv2_visualize_keypoints(frame, transformed_preds[idxs], num_pts=nParts, order=0)\n",
    "\n",
    "    cv2.imshow('frame', frame)\n",
    "\n",
    "    if is_part:\n",
    "        cv2.imwrite('demo/'+global_animal+'/'+img_path + '_all.jpg', np.concatenate( (frame, pred_mask), axis=1))\n",
    "        cv2.imwrite('demo/'+global_animal+'/'+img_path+'_seg.jpg', pred_mask)\n",
    "        # out.write(np.concatenate( (frame, pred_mask), axis=1))\n",
    "    cv2.imwrite('demo/'+global_animal+'/'+img_path, frame)\n",
    "\n",
    "    \n",
    "    if cv2.waitKey(1) & 0xFF == ord('q'):\n",
    "        break\n",
    "\n",
    "# Release everything if job is finished\n",
    "out.release()\n",
    "cv2.destroyAllWindows()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "python37",
   "language": "python",
   "name": "python37"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
